#include <rclcpp/rclcpp.hpp>
#include <sensor_msgs/msg/image.hpp>
#include <cv_bridge/cv_bridge.h>
#include <opencv2/opencv.hpp>
#include <apriltag/apriltag.h>
#include <apriltag/tag36h11.h>
#include <yaml-cpp/yaml.h>
#include <ament_index_cpp/get_package_share_directory.hpp>
#include <apriltag/apriltag_pose.h>
#include <tf2_ros/transform_broadcaster.h>
#include <geometry_msgs/msg/transform_stamped.hpp>
#include <tf2/LinearMath/Quaternion.h>
#include <tf2/LinearMath/Matrix3x3.h>
#include <std_msgs/msg/bool.hpp>
#include <unistd.h>
#include <sched.h>
#include <cstring>
#include <iomanip>
#include <sstream>
#include <tf2_geometry_msgs/tf2_geometry_msgs.hpp>

class AprilTagProcessor : public rclcpp::Node
{
public:
  AprilTagProcessor()
    : Node("apriltag_processor", 
          rclcpp::NodeOptions().automatically_declare_parameters_from_overrides(true)),
      apmControllerNameSpace_("/apm_drone")
  {
    std::this_thread::sleep_for(std::chrono::milliseconds(2000)); // 等待2秒


    // 检测时钟源并设置参数
    bool use_sim_time = check_clock_topic();
    this->set_parameter(rclcpp::Parameter("use_sim_time", use_sim_time));
    RCLCPP_INFO(this->get_logger(), "Using %s clock", 
               use_sim_time ? "simulation" : "system");

    // 初始化相机参数
    load_camera_config();

    // 初始化AprilTag检测器
    init_apriltag_detector();

    // 初始化TF广播器
    tf_broadcaster_ = std::make_unique<tf2_ros::TransformBroadcaster>(*this);

    // 创建订阅
    init_subscriptions();

    RCLCPP_INFO(this->get_logger(), "AprilTag Processor initialized");
  }

  ~AprilTagProcessor()
  {
    apriltag_detector_destroy(td_);
    tag36h11_destroy(tf_);
  }

private:
  // 检测时钟话题存在性
  bool check_clock_topic()
  {
    auto topics = this->get_topic_names_and_types();
    for (const auto &[name, types] : topics) {
      if (name == "/clock") return true;
    }
    return false;
  }

  void load_camera_config()
  {
    try {
      const std::string config_path = 
        ament_index_cpp::get_package_share_directory("move_to_gps_target_c") + 
        "/config/ost.yaml";
      
      YAML::Node config = YAML::LoadFile(config_path);
      
      // 加载相机矩阵
      std::vector<double> camera_data = config["camera_matrix"]["data"].as<std::vector<double>>();
      camera_matrix_ = cv::Mat(3, 3, CV_64F, camera_data.data()).clone();
      
      // 加载畸变系数
      std::vector<double> dist_data = config["distortion_coefficients"]["data"].as<std::vector<double>>();
      dist_coeffs_ = cv::Mat(1, 5, CV_64F, dist_data.data()).clone();

      // 提取内参
      fx_ = camera_matrix_.at<double>(0, 0);
      fy_ = camera_matrix_.at<double>(1, 1);
      cx_ = camera_matrix_.at<double>(0, 2);
      cy_ = camera_matrix_.at<double>(1, 2);

      RCLCPP_INFO(get_logger(), "Camera parameters loaded");
    } catch (const std::exception& e) {
      RCLCPP_FATAL(get_logger(), "Config loading failed: %s", e.what());
      rclcpp::shutdown();
    }
  }

  void init_apriltag_detector() {
    // 基础配置
    tf_ = tag36h11_create();
    td_ = apriltag_detector_create();
    apriltag_detector_add_family(td_, tf_);
    
    // 核心参数调整
    td_->quad_decimate = 1.0;      // 关闭降采样，保留细节
    td_->quad_sigma = 1.0;         // 适度高斯模糊，平滑阴影边缘
    td_->nthreads = 1;             // 多线程补偿计算量增加
    td_->refine_edges = 1;         // 启用边缘细化，修复断裂轮廓

    // 调整四边形检测阈值（通过td_->qtp参数）
    td_->qtp.min_white_black_diff = 10;  // 降低黑白最小差异（默认40）
    td_->qtp.max_line_fit_mse = 40.0;    // 放宽边缘拟合误差（默认10.0）
    td_->qtp.cos_critical_rad = 0.7;   // 允许更倾斜的四边形（默认0.9）
  }


  void init_subscriptions()
  {
    // 图像订阅
    sub_ = create_subscription<sensor_msgs::msg::Image>(
      "/camera",
      rclcpp::SensorDataQoS(),
      [this](const sensor_msgs::msg::Image::ConstSharedPtr& msg) {
        process_image(msg);
      });

    // 使能控制订阅
    enable_sub_ = create_subscription<std_msgs::msg::Bool>(
      apmControllerNameSpace_ + "/enable_apriltag_processor",
      10,
      [this](const std_msgs::msg::Bool::SharedPtr msg) {
        enable_apriltag_processor_ = msg->data;
        if (enable_apriltag_processor_) {
          last_process_time_ = this->get_clock()->now() - rclcpp::Duration(0, 500000000);
        }
      });
  }

  void process_image(const sensor_msgs::msg::Image::ConstSharedPtr& msg)
  {
    if (!enable_apriltag_processor_) return;

    const auto now = this->get_clock()->now();
    if ((now - last_process_time_).seconds() < 0.25) return;
    last_process_time_ = now;

    try {
      // 图像预处理
      cv::Mat image = cv_bridge::toCvCopy(msg, "bgr8")->image;
      cv::Mat undistorted, gray;
      cv::undistort(image, undistorted, camera_matrix_, dist_coeffs_);
      cv::cvtColor(undistorted, gray, cv::COLOR_BGR2GRAY);

cv::Ptr<cv::CLAHE> clahe = cv::createCLAHE();
clahe->setClipLimit(4.0);  // 调整对比度限制
cv::Mat enhanced_image;
clahe->apply(gray, enhanced_image);  // 应用CLAHE

      // AprilTag检测
      image_u8_t im = {enhanced_image.cols, enhanced_image.rows, enhanced_image.cols, enhanced_image.data};
      zarray_t* detections = apriltag_detector_detect(td_, &im);

      // 处理检测结果
      for (int i = 0; i < zarray_size(detections); i++) {
        apriltag_detection_t* det;
        zarray_get(detections, i, &det);

        // 姿态估计
        apriltag_pose_t pose;
        apriltag_detection_info_t info = {
          det, 2.0, fx_, fy_, cx_, cy_
        };
        estimate_tag_pose(&info, &pose);

        // 构建TF变换
        publish_transform(det->id, pose);
      }

      apriltag_detections_destroy(detections);
    } catch (const cv_bridge::Exception& e) {
      RCLCPP_ERROR(get_logger(), "Image processing error: %s", e.what());
    }
  }

  void publish_transform(int tag_id, const apriltag_pose_t& pose)
  {
    geometry_msgs::msg::TransformStamped transform;
    
    // 使用节点时钟获取时间戳
    transform.header.stamp = this->get_clock()->now();
    
    // 坐标系设置
    transform.header.frame_id = "camera_link";
    transform.child_frame_id = "tag_" + std::to_string(tag_id);

    // 设置平移（根据摄像头安装方向调整）
    transform.transform.translation.x = pose.t->data[2];  // Z轴
    transform.transform.translation.y = -pose.t->data[0]; // X轴反向
    transform.transform.translation.z = -pose.t->data[1]; // Y轴反向

    // 设置旋转
    tf2::Matrix3x3 rot_matrix(
      pose.R->data[0], pose.R->data[1], pose.R->data[2],
      pose.R->data[3], pose.R->data[4], pose.R->data[5],
      pose.R->data[6], pose.R->data[7], pose.R->data[8]
    );
    tf2::Quaternion q;
    rot_matrix.getRotation(q);
    q.normalize();
    transform.transform.rotation = tf2::toMsg(q);

    // 发布TF
    tf_broadcaster_->sendTransform(transform);

    // 打印调试信息
    // log_transform(tag_id, transform);
  }

  void log_transform(int tag_id, const geometry_msgs::msg::TransformStamped& transform)
  {
    std::ostringstream oss;
    oss << "Detected Tag " << tag_id << "\n"
        << "  Timestamp: " << transform.header.stamp.sec << "." 
        << std::setw(9) << std::setfill('0') << transform.header.stamp.nanosec << "\n"
        << "  Frame: " << transform.header.frame_id << " → " << transform.child_frame_id << "\n"
        << "  Translation: [" 
        << std::fixed << std::setprecision(3)
        << transform.transform.translation.x << ", "
        << transform.transform.translation.y << ", "
        << transform.transform.translation.z << "]\n"
        << "  Rotation: ["
        << transform.transform.rotation.x << ", "
        << transform.transform.rotation.y << ", "
        << transform.transform.rotation.z << ", "
        << transform.transform.rotation.w << "]";
    
    RCLCPP_INFO(this->get_logger(), "%s", oss.str().c_str());
  }

  // 成员变量
  bool enable_apriltag_processor_ = false;
  rclcpp::Time last_process_time_;
  std::string apmControllerNameSpace_;

  cv::Mat camera_matrix_;
  cv::Mat dist_coeffs_;
  double fx_, fy_, cx_, cy_;

  apriltag_family_t* tf_;
  apriltag_detector_t* td_;
  
  rclcpp::Subscription<sensor_msgs::msg::Image>::SharedPtr sub_;
  rclcpp::Subscription<std_msgs::msg::Bool>::SharedPtr enable_sub_;
  std::unique_ptr<tf2_ros::TransformBroadcaster> tf_broadcaster_;
};

bool bind_to_cpu(int cpu_core)
{
    cpu_set_t cpuset;
    CPU_ZERO(&cpuset);
    CPU_SET(cpu_core, &cpuset);
    pid_t pid = getpid();

    if (sched_setaffinity(pid, sizeof(cpu_set_t), &cpuset) == -1)
    {
        std::cerr << "Failed to set CPU affinity: " << strerror(errno) << std::endl;
        return false;
    }
    return true;
}

int main(int argc, char** argv)
{
    // 处理CPU绑定参数
    int cpu_core = -1;
    for (int i = 1; i < argc; ++i) {
        if (strcmp(argv[i], "--cpu") == 0 && i+1 < argc) {
            cpu_core = atoi(argv[i+1]);
            for(int j = i; j+2 <= argc; ++j)
                argv[j] = argv[j+2];
            argc -= 2;
            break;
        }
    }

    if(cpu_core != -1 && !bind_to_cpu(cpu_core)){
        std::cerr << "Continuing without CPU binding" << std::endl;
    }

    rclcpp::init(argc, argv);
    auto node = std::make_shared<AprilTagProcessor>();
    
    rclcpp::executors::SingleThreadedExecutor executor;
    executor.add_node(node);
    
    if(cpu_core != -1){
        cpu_set_t cpuset;
        CPU_ZERO(&cpuset);
        CPU_SET(cpu_core, &cpuset);
        if(pthread_setaffinity_np(pthread_self(), sizeof(cpu_set_t), &cpuset)){
            std::cerr << "Failed to set executor thread affinity" << std::endl;
        }
    }

    executor.spin();
    rclcpp::shutdown();
    return 0;
}
